import warnings

import numpy as np
import matplotlib.pyplot as plt
import time
import pygame


# 对一个数进行二进制编码，需要的位数
def get_encoded_bits(x):
    bits = 1
    while x > 1.99:
        x = x / 2
        bits += 1

    return bits


class Surviving(object):
    def __init__(self, n_agent=10, map_size=14, obs_dis=1, comm_dis=5, if_show_frame=False):
        super(Surviving, self).__init__()
        # 接受传入的参数
        self.n_agent = n_agent
        self.map_size = map_size
        self.obs_dis = obs_dis
        self.comm_dis = comm_dis
        self.if_show_frame = if_show_frame

        # pygame 配置
        if self.if_show_frame:
            pygame.init()
            self.size = (self.map_size * 20, self.map_size * 20)
            self.screen = pygame.display.set_mode(self.size)
            pygame.display.set_caption("Surviving")

        self.n_action = 5
        self.max_food = 6
        self.capability = 2 * self.n_agent
        self.resource_min = self.n_agent
        self.resource_max = int(self.n_agent * 1.5)
        self.pos_encoded_bits = get_encoded_bits(map_size + 1)

        self.maze = self.build_env()
        self.ants = []  # 智能体定义，随机位置
        for i in range(self.n_agent):
            self.ants.append([np.random.randint(0, map_size) + 1, np.random.randint(0, map_size) + 1])

        self.foods = []  # 定义每个智能体的食物量
        for i in range(self.n_agent):
            self.foods.append(self.max_food)

        self.n_resource = 8
        self.resource = []
        self.resource_pos = []  # 定义食物，随机位置
        for i in range(self.n_resource):
            self.resource_pos.append([np.random.randint(0, map_size) + 1, np.random.randint(0, map_size) + 1])
            self.resource.append(np.random.randint(self.resource_min, self.resource_max))

        self.steps = 0
        # 位置编码 + 资源编码 + 智能体编码 + 自身资源编码
        self.len_obs = self.pos_encoded_bits * 2 + 3 * 3 + 3 * 3 + 1
        self.detached_obs_one_length = self.pos_encoded_bits * 2 + 1

    # 判断位置是否在地图范围内
    def is_legal(self, x, y):
        return (x >= 1) & (x <= self.map_size) & (y >= 1) & (y <= self.map_size)

    def reset(self):

        self.maze = self.build_env()

        self.ants = []
        for i in range(self.n_agent):
            self.ants.append([np.random.randint(0, self.map_size) + 1, np.random.randint(0, self.map_size) + 1])

        self.foods = []
        for i in range(self.n_agent):
            self.foods.append(self.max_food)

        self.resource = []
        self.resource_pos = []
        for i in range(self.n_resource):
            self.resource_pos.append([np.random.randint(0, self.map_size) + 1, np.random.randint(0, self.map_size) + 1])
            self.resource.append(np.random.randint(self.resource_min, self.resource_max))

        return self.get_detached_obs(), self.get_adj()

    # 构建矩阵地图，设置边界
    def build_env(self):

        maze = np.zeros((self.map_size + 2, self.map_size + 2))
        for i in range(self.map_size + 2):
            maze[0][i] = -1
            maze[i][0] = -1
            maze[self.map_size + 1][i] = -1
            maze[i][self.map_size + 1] = -1

        return maze

    # 这部分的目的在于修改原有的智能体观测结果，从一个长的向量修改为三个短向量，用于分别处理
    # 三个向量分别为：1）智能体当前位置及食物情况 2）智能体相邻资源信息 3）智能体相邻智能体信息
    def get_detached_obs(self):
        obs = []
        maze_ant = np.zeros((self.map_size + 2, self.map_size + 2))
        for index in range(self.n_agent):
            x = self.ants[index][0]
            y = self.ants[index][1]
            maze_ant[x][y] = 1

        for index in range(self.n_agent):
            self_info = []
            resource_info = []
            agent_info = []

            x = self.ants[index][0]
            y = self.ants[index][1]
            for i in range(self.pos_encoded_bits):
                self_info.append(np.mod(x, 2))  # 取模 x % 2 对x编码
                x = int(x / 2)
            for i in range(self.pos_encoded_bits):
                self_info.append(np.mod(y, 2))  # 对y编码
                y = int(y / 2)
            self_info.append(self.foods[index])  # 食物信息

            x_t = self.ants[index][0]
            y_t = self.ants[index][1]
            # 获取自身3 * 3范围内资源
            for i in range(-self.obs_dis, self.obs_dis + 1):
                for j in range(-self.obs_dis, self.obs_dis + 1):
                    resource_info.append(self.maze[x_t + i][y_t + j])
            if len(resource_info) == self.detached_obs_one_length:
                pass
            elif len(resource_info) < self.detached_obs_one_length:
                for _ in range(self.detached_obs_one_length - len(resource_info)):
                    resource_info.append(0)

            # 获取自身3 * 3范围内智能体
            for i in range(-self.obs_dis, self.obs_dis + 1):
                for j in range(-self.obs_dis, self.obs_dis + 1):
                    agent_info.append(maze_ant[x_t + i][y_t + j])
            if len(agent_info) == self.detached_obs_one_length:
                pass
            elif len(agent_info) < self.detached_obs_one_length:
                for i in range(self.detached_obs_one_length - len(agent_info)):
                    agent_info.append(0)

            obs.append([self_info, resource_info, agent_info])

        return obs

    # 返回类型为List ，返回的shape为：[agent_num, obs_len]
    def get_obs(self):

        obs = []

        maze_ant = np.zeros((self.map_size + 2, self.map_size + 2))
        for index in range(self.n_agent):
            x = self.ants[index][0]
            y = self.ants[index][1]
            maze_ant[x][y] = 1

        # 对于每个智能体,有4个x编码，4个y编码。9个资源编码，9个智能体编码，1个自身食物编码.共27位
        for index in range(self.n_agent):
            h = []
            x = self.ants[index][0]
            y = self.ants[index][1]
            for i in range(self.pos_encoded_bits):
                h.append(np.mod(x, 2))  # 取模 x % 2 对x编码
                x = int(x / 2)
            for i in range(self.pos_encoded_bits):
                h.append(np.mod(y, 2))  # 对y编码
                y = int(y / 2)
            x_t = self.ants[index][0]
            y_t = self.ants[index][1]
            # 获取自身3 * 3范围内资源
            for i in range(-self.obs_dis, self.obs_dis + 1):
                for j in range(-self.obs_dis, self.obs_dis + 1):
                    h.append(self.maze[x_t + i][y_t + j])
            # 获取自身3 * 3范围内智能体
            for i in range(-self.obs_dis, self.obs_dis + 1):
                for j in range(-self.obs_dis, self.obs_dis + 1):
                    h.append(maze_ant[x_t + i][y_t + j])

            # 获取自己的食物状态
            h.append(self.foods[index])
            obs.append(h)

        return obs

    # 获得邻接矩阵，返回类型为ndarray，shape为(1, agent_num, agent_num)
    def get_adj(self):

        adj = np.zeros((1, self.n_agent, self.n_agent), dtype=np.int8)

        maze_ant = np.ones((self.map_size + 2, self.map_size + 2), dtype=np.int8) * -1
        # 获取所有智能体位置以及编号
        for index in range(self.n_agent):
            x = self.ants[index][0]
            y = self.ants[index][1]
            maze_ant[x][y] = index

        #
        for index in range(self.n_agent):
            x = self.ants[index][0]
            y = self.ants[index][1]

            # 在自身为中心，5 * 5范围内的智能体信息
            for i in range(-self.comm_dis, self.comm_dis + 1):
                for j in range(-self.comm_dis, self.comm_dis + 1):
                    if self.is_legal(x + i, y + j):
                        if maze_ant[x + i][y + j] != -1:  # 如果有智能体,邻接矩阵置1
                            adj[0][index][maze_ant[x + i][y + j]] = 1

        return adj

    def step(self, actions):
        """
        :param actions: 有五个值，分别对应左右上下和吃
        :return:
        """
        for i in range(self.n_agent):
            x = self.ants[i][0]
            y = self.ants[i][1]

            if actions[i] == 0:
                if self.maze[x - 1][y] != -1:
                    self.ants[i][0] = x - 1
            if actions[i] == 1:
                if self.maze[x + 1][y] != -1:
                    self.ants[i][0] = x + 1
            if actions[i] == 2:
                if self.maze[x][y - 1] != -1:
                    self.ants[i][1] = y - 1
            if actions[i] == 3:
                if self.maze[x][y + 1] != -1:
                    self.ants[i][1] = y + 1
            if actions[i] == 4:
                self.foods[i] += 2 * self.maze[x][y]
                self.maze[x][y] = 0

            # 食物不能大于max_food，不能小于0
            self.foods[i] = max(0, min(self.foods[i] - 1, self.max_food))

        # 没有食物就惩罚
        reward = [0.4] * self.n_agent
        for i in range(self.n_agent):
            if self.foods[i] == 0:
                reward[i] = - 0.2

        done = False
        # 环境的最大承载量，资源最多200点。大于环境承载量时，资源不新增
        if (self.maze.sum() + self.map_size * 4) > self.capability:
            return self.get_detached_obs(), self.get_adj(), reward, done

        for i in range(self.n_resource):

            # 资源点随机刷新资源
            x = self.resource_pos[i][0] + np.random.randint(-3, 4)
            y = self.resource_pos[i][1] + np.random.randint(-3, 4)

            if self.is_legal(x, y):

                # 资源大小随机，并且资源不大于5，消耗资源点的总资源
                num = np.random.randint(1, 4)
                self.maze[x][y] += num
                self.maze[x][y] = min(self.maze[x][y], 3)
                self.resource[i] -= num

                # 总资源小于0时，重新刷新资源点
                if self.resource[i] <= 0:
                    self.resource_pos[i][0] = np.random.randint(0, self.map_size) + 1
                    self.resource_pos[i][1] = np.random.randint(0, self.map_size) + 1
                    self.resource[i] = np.random.randint(self.resource_min, self.resource_max)

        return self.get_detached_obs(), self.get_adj(), reward, done

    # 可视化地图，现有的地图中： -1代表边界，0代表正常地面
    def render(self):
        data = self.maze
        for i in range(self.n_agent):
            x = self.ants[i][0]
            y = self.ants[i][1]
            data[x][y] = 1

        for i in range(self.n_resource):
            x = self.resource_pos[i][0]
            y = self.resource_pos[i][1]
            data[x][y] = 2

        fig, ax = plt.subplots()
        ax.imshow(data, extent=(0, data.shape[1], 0, data.shape[0]))
        for i in range(self.n_agent):
            x = self.ants[i][0]
            y = self.ants[i][1]
            food = self.foods[i]
            ax.text(y, data.shape[0] - x - 1, 'A{0}_{1}'.format(i, food), ha="left", va="bottom", color="w",
                    fontsize='xx-small')
        for i in range(self.n_resource):
            x = self.resource_pos[i][0]
            y = self.resource_pos[i][1]
            res = self.resource[i]
            ax.text(y, data.shape[0] - x - 1, 'R{0}_{1}'.format(i, res), ha="left", va="bottom", color="black",
                    fontsize='xx-small')
        ax.grid()
        ax.set_xticks(np.arange(data.shape[1]))
        ax.set_yticks(np.arange(data.shape[0]))
        # fig.tight_layout()
        plt.show()
        time.sleep(0.05)

    def show_frame(self):
        if not self.if_show_frame:
            return

        self.screen.fill((255, 255, 255))

        grid_surface = pygame.Surface(self.size)
        grid_surface.fill((255, 255, 255))

        for x in range(0, self.size[0], 20):
            pygame.draw.line(grid_surface, (0, 0, 0), (x, 0), (x, self.size[1]))
        for y in range(0, self.size[1], 20):
            pygame.draw.line(grid_surface, (0, 0, 0), (0, y), (self.size[0], y))

        # 画智能体
        for i in range(self.n_agent):
            pygame.draw.circle(grid_surface, (0, 0, 255), (self.ants[i][0] * 20 + 10, self.ants[i][1] * 20 + 10), 5)

        # 画资源点
        for i in range(self.n_resource):
            pygame.draw.circle(grid_surface, (0, 255, 0),
                               (self.resource_pos[i][0] * 20 + 10, self.resource_pos[i][1] * 20 + 10), 5)

        self.screen.blit(grid_surface, (0, 0))
        pygame.display.flip()
        time.sleep(0.1)
